import os
import time
import cv2
import numpy as np
import tensorflow as tf
from model.model_head_CBAM import efficientnetv2_m as create_model

# === 参数 ===
input_size = (384, 384)
num_classes = 15
label_names = [
    'colorResult_grey', 'colorResult_white', 'colorResult_yellow',
    'shapeResult_ToothMarks', 'shapeResult_fat', 'shapeResult_normal',
    'shapeResult_thin', 'textureResult_dark', 'textureResult_normal',
    'textureResult_tender', 'textureResult_water', 'thicknessResult_Stripping',
    'thicknessResult_ecchymosis', 'thicknessResult_greasy', 'thicknessResult_thin'
]

# === 1. 配置多 GPU ===
# 列出所有可见 GPU，并打开 memory growth，避免一次性占满显存
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# 创建 MirroredStrategy，会自动使用所有可见 GPU
strategy = tf.distribute.MirroredStrategy()
print(f"Number of devices: {strategy.num_replicas_in_sync}")

# === 2. 在 strategy.scope() 下创建并加载模型 ===

with strategy.scope():
    model = create_model(num_classes=num_classes)
    model.build((None, *input_size, 3))
    model.load_weights("./efficientnetv2_best_head_CBAM.h5")


# === 3. 预处理函数（同前） ===
def preprocess_image(img_path):
    img = cv2.imread(img_path)
    if img is None:
        raise FileNotFoundError(f"图像不存在：{img_path}")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, input_size)
    img = img.astype(np.float32) / 255.0
    return img  # 注意：这里去掉了 batch 维度


# === 4. 多 GPU 推理函数 ===
def predict_images_multi_gpu(img_paths):
    # 1) 读入并预处理，堆叠成 batch
    imgs = [preprocess_image(p) for p in img_paths]
    batch = np.stack(imgs, axis=0)  # shape: (batch_size, H, W, C)

    # 2) 确保 log 目录存在
    logdir = "./logs"
    os.makedirs(logdir, exist_ok=True)

    # 3) 同时启动 profiler 和计时
    tf.profiler.experimental.start(logdir)
    t0 = time.time()

    # 4) 直接 predict，Keras 会按 replica 平分 batch 并行执行
    preds = model.predict(batch)

    tf.profiler.experimental.stop()
    t1 = time.time()

    # 5) 输出结果
    print(f"\nBatch size = {len(img_paths)}  （共用 {strategy.num_replicas_in_sync} 个 GPU）")
    print(f"Batch 推理总耗时：{(t1 - t0):.4f} 秒，平均每张：{(t1 - t0) / len(img_paths):.4f} 秒")
    print(f"性能分析日志已保存至：{logdir}")


# === 5. 示例调用 ===
predict_images_multi_gpu([r"valid/1.png" for i in range(0, 200)])